﻿Imports System.IO
Imports System.Threading
Imports System.ServiceProcess
Imports System.Runtime.Remoting.Messaging
Imports System.Configuration.ConfigurationManager
Imports Slythe.FileSystemProfiler.ProfilerService.Models



Public Class ProfilerService



#Region "Private Members"



    Private _controllerThread As Thread


    Private _profilerThreads As Dictionary(Of Thread, ProfilerThreadStateEnum)


    Private _abort As Boolean



#End Region



#Region "Enums"



    Private Enum ProfilerThreadStateEnum As Integer
        UNKNOWN = 0
        RUNNING = 1
        FINISHED = 2
    End Enum



#End Region



#Region "Constructors"



    <MTAThread()> _
    Shared Sub Main()


        Dim serviceToRun As ProfilerService = New ProfilerService


        If Environment.UserInteractive Then


            serviceToRun.OnStart(Nothing)

            Console.WriteLine("Press any key to stop the service")
            Console.ReadKey(True)

            serviceToRun.OnStop()


        Else


            ServiceBase.Run(serviceToRun)


        End If


    End Sub



    Protected Overrides Sub OnStart(ByVal args() As String)
        ' Add code here to start your service. This method should set things
        ' in motion so your service can do its work.

        _controllerThread = New Thread(AddressOf ProfileDrives)

        With _controllerThread
            .Name = "Control Thread"
            .Priority = ThreadPriority.BelowNormal
        End With

        _controllerThread.Start()

    End Sub



#End Region



#Region "Delegates"



    Delegate Function dlgDriveProfiler(ByVal drive As DriveInfo) As Integer



#End Region



#Region "Stop/Tear-down"



    Protected Overrides Sub OnStop()
        ' Add code here to perform any tear-down necessary to stop your service.

        'Kill threads
        If Not _profilerThreads Is Nothing Then

            KillProfilerThreads()

        End If

    End Sub



    Private Sub KillProfilerThreads()


        Dim selectedThread As Thread = Nothing

        For Each profiler As KeyValuePair(Of Thread, ProfilerThreadStateEnum) In _profilerThreads

            selectedThread = profiler.Key

            selectedThread.Abort()

        Next


    End Sub



#End Region





    Private Sub ProfileDrives()


        Dim drivesToProfile As List(Of DriveInfo) = Nothing


        Do Until _abort


            'Get drives to profile
            drivesToProfile = getDrivesToProfile()


            'create a thread per drive and start profiling
            startThreadsForDriveProfiling(drivesToProfile)


            'Check that all the threads are completed

            If AreAllProfilerThreadsFinished() Then

                'Sleep for a long time (until the next run)
                Thread.Sleep(getSleepDuration)

            Else

                'Sleep for a second then re-evaluate
                Thread.Sleep(TimeSpan.FromSeconds(1))

            End If


        Loop



    End Sub



    Private Sub startThreadsForDriveProfiling(ByVal drives As List(Of DriveInfo))


        Dim dlgProfiler As dlgDriveProfiler


        For Each drive As DriveInfo In drives


            dlgProfiler = New dlgDriveProfiler(AddressOf ProfileDrive)


            dlgProfiler.BeginInvoke(drive, _
                                    New AsyncCallback(AddressOf driveProfileComplete), _
                                    Nothing)


        Next


    End Sub



    Private Function ProfileDrive(ByVal drive As DriveInfo) As ProfilerThreadStateEnum


        Dim threadState As ProfilerThreadStateEnum = _
            ProfilerThreadStateEnum.RUNNING


        NameThreadAfterDrive(drive)

        AddThreadToCollection()


        Dim files As New List(Of FileInfo)
        Dim directories As New List(Of DirectoryInfo)


        Dim driveContentsInfo As New DriveContentsInformationModel

        driveContentsInfo = recurseDirectories(driveContentsInfo, _
                                               drive.RootDirectory)


        Dim serialiser As Serialisation.SerialiseDriveStructure = _
            New Serialisation.SerialiseDriveStructure(driveContentsInfo)

        serialiser.Start()


        threadState = ProfilerThreadStateEnum.FINISHED


        UpdateThreadStatus(threadState)


        Return threadState


    End Function



    Private Function recurseDirectories(ByVal processedItems As DriveContentsInformationModel, _
                                        ByVal directoryToExamine As DirectoryInfo) As DriveContentsInformationModel

 

        For Each givenDirectory In directoryToExamine.GetDirectories


            Try

                processedItems.Files.AddRange(givenDirectory.GetFiles)

                processedItems.Directories.Add(givenDirectory)

                processedItems = recurseDirectories(processedItems, _
                                                    givenDirectory)


            Catch ex As UnauthorizedAccessException

                Debug.WriteLine("Access denied on directory " & givenDirectory.FullName)

            End Try


        Next


        Return processedItems


    End Function



    Private Function getDrivesToProfile() As List(Of DriveInfo)


        Dim driveLettersToProfile As New List(Of String)
        driveLettersToProfile.AddRange(AppSettings("Drives.DrivesToProfile").ToString.Split(","))


        Dim drivesToProfile As New List(Of DriveInfo)

        Dim driveInfo As DriveInfo = Nothing
        Dim dirInfo As DirectoryInfo = Nothing

        For Each driveValue As String In driveLettersToProfile


            dirInfo = New DirectoryInfo(driveValue)

            If dirInfo.Exists Then

                driveInfo = New DriveInfo(dirInfo.Root.FullName)

                drivesToProfile.Add(driveInfo)

            End If


        Next


        Return drivesToProfile


    End Function



    Private Sub driveProfileComplete(ByVal result As IAsyncResult)


        Dim resultClass As AsyncResult = CType(result,  _
                                               AsyncResult)

        Dim dlgDriveProfiler As dlgDriveProfiler = CType(resultClass.AsyncDelegate,  _
                                                         dlgDriveProfiler)

        Dim threadState As ProfilerThreadStateEnum = _
            ProfilerThreadStateEnum.UNKNOWN


        threadState = dlgDriveProfiler.EndInvoke(result)

        Console.WriteLine("Thread completed with state: " & threadState.ToString)


    End Sub



    Private Sub NameThreadAfterDrive(ByVal drive As DriveInfo)


        Dim threadName As String = String.Empty


        If String.IsNullOrWhiteSpace(drive.VolumeLabel) Then

            threadName = drive.Name

            threadName = threadName.Replace(Path.VolumeSeparatorChar, String.Empty)
            threadName = threadName.Replace(Path.PathSeparator, String.Empty)
            threadName = threadName.Replace(Path.AltDirectorySeparatorChar, String.Empty)


        Else

            threadName = drive.VolumeLabel.Replace(" ", "")

        End If


        Thread.CurrentThread.Name = threadName


    End Sub



    Private Sub AddThreadToCollection()


        'make sure its not null
        If _profilerThreads Is Nothing Then

            _profilerThreads = New Dictionary(Of Thread, ProfilerThreadStateEnum)

        End If


        If Not _profilerThreads.ContainsKey(Thread.CurrentThread) Then

            _profilerThreads.Add(Thread.CurrentThread, ProfilerThreadStateEnum.RUNNING)

        Else

            _profilerThreads.Item(Thread.CurrentThread) = ProfilerThreadStateEnum.RUNNING

        End If



    End Sub



    Private Sub UpdateThreadStatus(ByVal threadStatus As ProfilerThreadStateEnum)


        'make sure its not null
        If _profilerThreads Is Nothing Then

            'TODO: This shouldn't be the case, and if it is there has been some sort of error.

        End If


        If Not _profilerThreads.ContainsKey(Thread.CurrentThread) Then

            'TODO: Again this shouldn't be the case, if the thread isn't in the collection then there has been an error

        Else

            _profilerThreads.Item(Thread.CurrentThread) = threadStatus

        End If


    End Sub



    Private Function getSleepDuration() As Integer


        'set to a default value of 24 hours
        Dim sleepDurationInHours As Decimal = 24
        Dim sleepDurationInMilliseconds As Integer = Integer.MinValue

        Dim configValue As String = AppSettings("Profiling.FrequencyToProfile")


        'try the conversion
        If Not Decimal.TryParse(configValue, sleepDurationInHours) Then

            'TODO: Raise the error that the config value is wrong

        End If


        sleepDurationInMilliseconds = TimeSpan.FromHours(sleepDurationInHours).TotalMilliseconds


        Return sleepDurationInMilliseconds


    End Function



    Private Function AreAllProfilerThreadsFinished() As Boolean


        Dim allFinished As Boolean = True


        'evaluate the collection
        If Not _profilerThreads Is Nothing Then

            'TODO: Make this collection a thread safe collection to prevent errors when evaluating
            For Each profilingThread As KeyValuePair(Of Thread, ProfilerThreadStateEnum) In _profilerThreads

                If profilingThread.Value = ProfilerThreadStateEnum.RUNNING _
                    Or profilingThread.Value = ProfilerThreadStateEnum.UNKNOWN Then

                    'If anything is still running or in unknown state then it is not all complete
                    allFinished = False

                    Exit For

                End If

            Next

        Else

            'If the collection is not yet instantiated then it can't be finished
            allFinished = False

        End If


        Return allFinished


    End Function



End Class